package net.xuele.learn.flink.book;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

/**
 * @Author patrick
 * @Date 2023/7/20 18:18
 * @Description
 */
public class ResettableCountSource implements SourceFunction<Long>, CheckpointedFunction {

    private boolean isRunning = true;

    private Long offset = 0L;

    private ListState<Long> offsetState;

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        offsetState.clear();
        offsetState.add(offset);
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        ListStateDescriptor<Long> desc = new ListStateDescriptor<>("offset", Long.class);
        offsetState = context.getOperatorStateStore().getListState(desc);

        for (Long l : offsetState.get()) {
            offset = l;
        }
    }

    @Override
    public void run(SourceContext<Long> ctx) throws Exception {
        final Object checkpointLock = ctx.getCheckpointLock();
        while (isRunning && offset < Long.MAX_VALUE) {
            synchronized (checkpointLock) {
                offset += 1;
                ctx.collect(offset);
            }
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
    }
}
